import Sugarscape,Agent
import numpy
import sys
import matplotlib.pyplot as plt
import csv

"""
A class to handle data collection for a Sugarscape object
"""

class DataCollect:
    def __init__(self,sugarscape):
        self.sugarscape=sugarscape
        self.wealths=[]
        self.Analyze()
        
    def Analyze(self):
        for agent in self.sugarscape.agents:
            self.wealths.append(agent.sugar_reserve)
        self.sd=numpy.std(self.wealths)
        self.avg=numpy.average(self.wealths)
        self.med=numpy.median(self.wealths)
        self.total_wealth=sum(self.wealths)
        self.num_agents=len(self.sugarscape.agents)
        self.maximum = max(self.wealths)
        self.minimum = min(self.wealths)
        self.rnge = self.maximum - self.minimum
        self.gini = self.calc_gini(self.wealths)
        self.bottom_quartile = self.calc_bottom_quartile(self.wealths)
        

    def Pack(self):
    #return a packed list describing properties of the sugarscape
        return [self.sugarscape.tax_rate*100,self.sd,self.avg,self.med,self.maximum,self.minimum,self.rnge,self.total_wealth,self.gini, self.bottom_quartile]

    def calc_bottom_quartile(self, wealths):
       sortWealths=sorted(wealths)
       return sum(sortWealths[:len(sortWealths)/4])/(len(sortWealths)/4.0) #average the bottom quartile

    def calc_gini(self, wealths):
        """Return computed Gini coefficient.
        :note: follows basic formula
        :see: `calc_gini2`
        :contact: aisaac AT american.edu
        """
        x = sorted(wealths)  # increasing order
        N = len(x)
        B = sum( xi * (N-i) for i,xi in enumerate(x) ) / (N*sum(x))
        return 1 + (1./N) - 2*B

    def make_wealth_histogram(self, scale):
        '''makes a histogram with different
        bins'''
        histogram = {}
        for wealth in self.wealths:
            wealth = int(wealth/scale)
            wealth_bin = wealth*scale
            histogram[wealth_bin] = histogram.get(wealth_bin, 0) + 1
        return histogram

    def pmf(self, tax_rate, scale=40):
        hist = self.make_wealth_histogram(scale)
        plt.bar(hist.keys(), hist.values(), width=scale, color='0.85')
        plt.xlabel('wealth', fontsize=15)
        plt.ylabel('number of agents', fontsize=15)
        plt.title('Wealth distribution, ' + str(tax_rate*100) + '% tax rate', fontsize=20)
        plt.show()
    
    
    def make_data_file(self, filename):
	    datafile = open(filename + ".csv", 'w')
	    datafile.write("tax_rate, standard deviation, average, median, maximum, minimum, range, total_wealth, gini\n")
	    tax_rates = [x * 0.001 for x in range(0, 250)]
	    medians = []
	    std_devs = []
	    total_wealths = []
	    ginis = []
	    bottoms = []
	    a = Sugarscape.Sugarscape(tax_rate=0)
	    for i in range(300):
	        a.nextstep()
	    d = DataCollect(a)
	    #~ d.pmf()
	    for rate in tax_rates:
	        a = Sugarscape.Sugarscape(tax_rate=rate)
	        i = 0
	        print rate
	        while i <= 300:
	            a.nextstep()
	            if i == 300:
	                d = DataCollect(a)
	                data = d.Pack()
	                medians.append(data[3])
	                std_devs.append(data[1])
	                total_wealths.append(data[7])
	                ginis.append(data[8])
	                bottoms.append(data[9])
	                datafile.write(",".join(map(str, data)) + "\n")
	            i += 1
	    return data
	    
	    #~ datafile.close()
	    #~ plt.figure(1)
	    #~ plt.subplot(221)
	    #~ plt.plot(tax_rates, ginis)
	    #~ plt.xlabel('Tax Rates')
	    #~ plt.ylabel('Gini coefficient')
	    #~ plt.subplot(222)
	    #~ plt.plot(tax_rates, std_devs)
	    #~ plt.xlabel('Tax Rates')
	    #~ plt.ylabel('Standard deviation')
	    #~ plt.subplot(223)
	    #~ plt.plot(tax_rates, medians)
	    #~ plt.xlabel('Tax Rates')
	    #~ plt.ylabel('Median wealth')
	    #~ plt.subplot(224)
	    #~ plt.plot(tax_rates, total_wealths)
	    #~ plt.xlabel('Tax Rates')
	    #~ plt.ylabel('Total wealth')
	    #~ plt.show()
	    
    def __str__(self):
        return "TAX RATE: " + str(self.sugarscape.tax_rate*100) + "%\n------\n  average: " + str(self.avg) + "\n  standard deviation: " + str(self.sd) + "\n  median: " + str(self.med) + "\n  maximum: " + str(self.maximum) + "\n  minimum: " + str(self.minimum) + "\n  range: " + str(self.rnge) + "\n  total wealth: " + str(self.total_wealth) + "\n  Number of agents: " + str(self.num_agents) + "\n  Gini: " + str(self.gini)

if __name__=='__main__':
#	pass
	#~ 
#~ <<<<<<< HEAD
   # a = Sugarscape.Sugarscape(tax_rate=0)
    #a = DataCollect(a)
   # print a.calc_bottom_quartile(range(100))
    #~ for i in range(300):
        #~ a.nextstep()
    #~ d = DataCollect(a)
    #~ d.Pack()
    #~ plt.plot(d.wealths,'o')
    #~ 
    #~ plt.show()
    #~ 
#~ 
    #~ datafile = open('tax_data_dan.csv', 'w')
    #~ datafile.write("tax_rate, standard deviation, average, median, maximum, minimum, range, total_wealth, gini\n")
    #~ tax_rates = [x * 0.001 for x in range(0, 250)]
    #~ medians = []
    #~ std_devs = []
    #~ total_wealths = []
    #~ ginis = []
    #~ for rate in tax_rates:
        #~ a = Sugarscape.Sugarscape(tax_rate=rate)
        #~ i = 0
        #~ print rate
        #~ while i <= 300:
            #~ a.nextstep()
            #~ if i == 300:
                #~ d = DataCollect(a)
                #~ print d
                #~ data = d.Pack()
                #~ medians.append(data[3])
                #~ std_devs.append(data[1])
                #~ total_wealths.append(data[7])
                #~ ginis.append(data[8])
                #~ datafile.write(",".join(map(str, data)) + "\n")
            #~ i += 1
    #~ datafile.close()
#~ =======
     a = Sugarscape.Sugarscape(tax_rate=.2)
     for i in range(300):
         a.nextstep()
     d = DataCollect(a)
     d.pmf(.2, scale=5)
     #~ 
#~ >>>>>>> d1de633458ec477da0d19e2cdd7b7975445ea226
#~ 
 #    datafile = open('tax_data_changedparam.csv', 'w')
 #    datafile.write("tax_rate,standard deviation,average,median,maximum,minimum,range,total_wealth,gini\n")
 #    tax_rates = [x * 0.001 for x in range(0, 250)]
 #    medians = []
 #    std_devs = []
 #    total_wealths = []
 #    ginis = []
 #    for rate in tax_rates:
 #        a = Sugarscape.Sugarscape(tax_rate=rate)
 #        i = 0
 #        print rate
 #        while i <= 300:
 #            a.nextstep()
 #            if i == 300:
 #                d = DataCollect(a)
 #                print d
 #                data = d.Pack()
 #                medians.append(data[3])
 #                std_devs.append(data[1])
 #                total_wealths.append(data[7])
 #                ginis.append(data[8])
 #                datafile.write(",".join(map(str, data)) + "\n")
 #            i += 1
 #    datafile.close()
 
    #~ plt.figure(1)
    #~ plt.subplot(221)
    #~ plt.plot(tax_rates, ginis)
    #~ plt.xlabel('Tax Rates')
    #~ plt.ylabel('Gini coefficient')
    #~ plt.subplot(222)
    #~ plt.plot(tax_rates, std_devs)
    #~ plt.xlabel('Tax Rates')
    #~ plt.ylabel('Standard deviation')
    #~ plt.subplot(223)
    #~ plt.plot(tax_rates, medians)
    #~ plt.xlabel('Tax Rates')
    #~ plt.ylabel('Median wealth')
    #~ plt.subplot(224)
    #~ plt.plot(tax_rates, total_wealths)
    #~ plt.xlabel('Tax Rates')
    #~ plt.ylabel('Total wealth')
    #~ plt.show()
#~ 
